### Common functions to deal with the output time and land_units
def open_files (DataSource, SceName, TabName, CaseName, VarName):
# read files from three different data source, ncar(cesm2-waccm), esgf, and e3sm
# return xarray dataset
import os
if DataSource == 'ncar':
TopDir = dir_cesm2
if TabName == 'Lmon':
FileTemp = TopDir + "/{}/land_monthly/{}.clm2.h0.{}.204001-210012.nc".format(CaseName, CaseName, VarName)
elif TabName == 'Amon':
FileTemp = TopDir + "/{}/atmos_monthly/{}.cam.h0.{}.204001-210012.nc".format(CaseName, CaseName, VarName)
if not os.path.exists(FileTemp):
FileTemp = FileTemp.replace('204001-210012', '*')
if DataSource == 'esgf':
TopDir = dir_esgf
FileTemp = TopDir + "{}/r1i1p1f1/{}/{}/gn/*/*.nc".format(SceName, TabName, VarName)
FileList = glob.glob(FileTemp)
filsel=None
versel=0
for fl in FileList:
tmparr = fl.split('/')
if int(tmparr[tmparr.index('gn')+1][1:]) >= versel:
filsel = tmparr[tmparr.index('gn')+1]
versel = int(tmparr[tmparr.index('gn')+1][1:])
if filsel:
FileTemp = TopDir + "{}/r1i1p1f1/{}/{}/gn/{}/*.nc".format(SceName, TabName, VarName, filsel)
if DataSource == 'e3sm':
TopDir = Dir_Geo_E3SM
if TabName == 'Lmon':
FileTemp = TopDir + "{}/rgr/lnd/{}_*.nc".format(CaseName, VarName)
elif TabName == 'Amon':
FileTemp = TopDir + "{}/rgr/atm/{}_*.nc".format(CaseName, VarName)
elif TabName == 'fx':
FileTemp = TopDir + "{}/rgr/fx/{}_*.nc".format(CaseName, VarName)
print (FileTemp)
if '*' in FileTemp:
ds = xr.open_mfdataset(FileTemp)
else:
ds = xr.open_dataset(FileTemp)
return ds
def gm_by_area_wght(ds, vname, OpType, LandInfo, area, landfrac, TabName, TimeShift=True):
# get the global mean weighted by area, ds is the xarray data array.
# time shift is needed for the raw outputs of cesm and e3sm, as time in time array is
# generally is the time at the end of integration for dumping history files.
# LandInfo: true, the xarray data array or data set consisting of the area and land fraction.
if TimeShift:
if TabName == 'Lmon':
ds["time0"] = ds.time_bounds[:,0].dt.ceil('D')
if TabName == 'Amon':
try:
ds["time0"] = ds.time_bnds[:,0]
except:
ds["time0"] = ds.time_bounds[:,0].dt.ceil('D')
dd = ds.set_index(time="time0")
else:
dd = ds
dayspermon=[31, 28, 31, 30, 31, 30, 31, 31, 30, 31, 30, 31]
dpm = np.zeros(dd.time.size)
for j in range (dd.time.size):
k = j%12
dpm[j] = dayspermon[k]
if LandInfo:
if TabName == 'Amon':
xarea = dd.areacella
else:
xarea = dd.area
else:
xarea = area.areacella
if LandInfo:
if TabName == 'Amon':
xlandfrac = dd.sftlf
else:
xlandfrac = dd.landfrac
else:
xlandfrac = landfrac.sftlf
xrdpm = xr.DataArray(data=dpm,
dims=["time"],
coords=dict(time=dd.time.values,),
attrs=dict(description="days per months", units="days",),)
# if the LandInfo == True, the area and fraction data are from the ds, so will be no problem, but when it is false, it will cause problem
if OpType == 'sum':
if TabName == 'Amon':
TsGBL = (dd[vname] * xarea.values[None,:,:] * 1.e6 * 86400 * xrdpm * 1.e-15).sum(dim=("lat", "lon")).groupby('time.year').sum('time')
else:
if LandInfo:
TsGBL = (dd[vname] * xarea * 1.e6 * xlandfrac * 86400 * xrdpm * 1.e-15).sum(dim=("lat", "lon")).groupby('time.year').sum('time')
else:
TsGBL = (dd[vname] * xarea.values * 1.e6 * xlandfrac.values * 86400 * xrdpm * 1.e-15).sum(dim=("lat", "lon")).groupby('time.year').sum('time')
elif OpType == 'avg':
if LandInfo:
if TabName == 'Amon':
temp = (dd[vname] * xarea * 1.e6 * xrdpm).sum(dim=("lat", "lon")) / (xarea * 1.e6 ).sum(dim=("lat", "lon"), skipna=True)
else:
temp = (dd[vname] * xarea * 1.e6 * xlandfrac * xrdpm).sum(dim=("lat", "lon")) / (xarea * 1.e6 * xlandfrac).sum(dim=("lat", "lon"), skipna=True)
else:
if TabName == 'Amon':
temp = (dd[vname] * xarea.values[None,:,:] * 1.e6 * xrdpm).sum(dim=("lat", "lon")) / (xarea * 1.e6 ).sum(dim=("lat", "lon"), skipna=True)
else:
temp = (dd[vname] * xarea.values[None,:,:] * 1.e6 * xlandfrac.values[None,:,:] * xrdpm).sum(dim=("lat", "lon")) / (xarea * 1.e6 * xlandfrac.values).sum(dim=("lat", "lon"), skipna=True)
TsGBL = temp.groupby('time.year').sum('time') / 365.
return TsGBL
def plot_ts_diff(VarList, CumDiff, LabList, Ylabel, VarE3SM=[]):
if CumDiff:
fig, faxs = plt.subplots(1,3)
else:
fig, faxs = plt.subplots(1,3)
axs = faxs.flatten()
#ax2 = ax.twinx()
LnColors_1 = ['#17becf', '#17becf', '#ff7f0e', '#ff7f0e', 'k']
for i, var in enumerate(VarList):
var.plot(ax=axs[0], label=LabList[i], color=LnColors_1[i])
LnColors_2 = ['#9467bd', 'b']
if VarE3SM:
e3smlab = ['E3SMSC-Geo-2.0', 'E3SMSC-SSP534OS']
for i, var in enumerate(VarE3SM):
if 'nbp' in Ylabel:
var[1:].plot(ax=axs[0], label=e3smlab[i], color=LnColors_2[i])
else:
var.plot(ax=axs[0], label=e3smlab[i], color=LnColors_2[i])
for i, var in enumerate(VarList[:-1]):
if CumDiff:
(var-VarList[-1]).plot.line('--', color=LnColors_1[i], linewidth=3, ax=axs[1])
(var-VarList[-1]).cumsum().plot.line('--', label=LabList[i], color=LnColors_1[i], linewidth=3, ax=axs[2])
else:
(var-VarList[-1]).plot.line('--', color=LnColors_1[i], linewidth=3, ax=axs[1])
if VarE3SM:
e3smlab = ['E3SMSC-Geo-2.0', 'E3SMSC-SSP534OS']
for i, var in enumerate(VarE3SM[:-1]):
if CumDiff:
if 'nbp' in Ylabel:
(var-VarE3SM[-1])[1:].cumsum().plot.line('--', label=e3smlab[i], color=LnColors_2[i], linewidth=3, ax=axs[2])
(var-VarE3SM[-1])[1:].plot.line('--', color=LnColors_2[i], linewidth=3, ax=axs[1])
else:
(var-VarE3SM[-1]).cumsum().plot.line('--', label=e3smlab[i], color=LnColors_2[i], linewidth=3, ax=axs[2])
(var-VarE3SM[-1]).plot.line('--', color=LnColors_2[i], linewidth=3, ax=axs[1])
else:
(var-VarE3SM[-1]).plot.line('--', color=LnColors_2[i], linewidth=3, ax=axs[1])
#ax2.tick_params(axis='y', labelcolor='blue')
#ax2.axhline(y=0.0, color='b', linestyle='--')
axs[0].legend(prop={'size': 12})
axs[2].legend(prop={'size': 12})
axs[0].set_ylabel(Ylabel, fontsize=20)
if CumDiff:
axs[1].set_ylabel(Ylabel+'Diff', fontsize=20)
axs[2].set_ylabel('cum'+Ylabel+'Diff', fontsize=20)
else:
axs[1].set_ylabel(Ylabel+'Diff', fontsize=20)
axs[2].axis('off')
def plot_ts(SecName, TabName, OpType, Vname_NCAR, Vname_ESGF, Sca_NCAR, Sca_ESGF, Cumdiff, Ylabel, noLimit = False, Vname_E3SM=None, Sca_E3SM=None, Fct_NCAR=0.0, Fct_ESGF=0.0, Fct_E3SM=0.0, SkipPlot=False):
opstr = ['+', '-', '*', '/']
## NCAR
if any (e in Vname_NCAR for e in opstr):
# sympy
#StrVname_NCAR.translate({ord(c):' ' for c in '+-*/'})
Sym_Expr_NCAR = sympify(Vname_NCAR)
FreeSymb_NCAR = list(Sym_Expr_NCAR.free_symbols)
fNCAR = lambdify(FreeSymb_NCAR, Sym_Expr_NCAR, "numpy")
dsTmp15_1 = []
dsTmp15_2 = []
dsTmp20_1 = []
dsTmp20_2 = []
for s in FreeSymb_NCAR:
dsTmp15_1.append(open_files('ncar', SecName.upper()+'-PK', TabName, GeoPK15_case1, s.name))
dsTmp20_1.append(open_files('ncar', SecName.upper()+'-PK', TabName, GeoPK20_case1, s.name))
dsTmp15_2.append(open_files('ncar', SecName.upper()+'-PK', TabName, GeoPK15_case2, s.name))
dsTmp20_2.append(open_files('ncar', SecName.upper()+'-PK', TabName, GeoPK20_case2, s.name))
dsVar15_1 = dsTmp15_1[0].copy()
dsVar15_1["uniqName"] = fNCAR(*[ds[s.name] for s, ds in zip(FreeSymb_NCAR,dsTmp15_1)])
dsVar15_2 = dsTmp15_2[0].copy()
dsVar15_2["uniqName"] = fNCAR(*[ds[s.name] for s, ds in zip(FreeSymb_NCAR,dsTmp15_2)])
dsVar20_1 = dsTmp20_1[0].copy()
dsVar20_1["uniqName"] = fNCAR(*[ds[s.name] for s, ds in zip(FreeSymb_NCAR,dsTmp20_1)])
dsVar20_2 = dsTmp20_2[0].copy()
dsVar20_2["uniqName"] = fNCAR(*[ds[s.name] for s, ds in zip(FreeSymb_NCAR,dsTmp20_2)])
if not noLimit:
dsVar15_1["uniqName"] = dsVar15_1["uniqName"].where(np.logical_and(dsVar15_1["uniqName"] > 0, dsVar15_1["uniqName"] < 1))
dsVar20_1["uniqName"] = dsVar20_1["uniqName"].where(np.logical_and(dsVar20_1["uniqName"] > 0, dsVar20_1["uniqName"] < 1))
dsVar15_2["uniqName"] = dsVar15_2["uniqName"].where(np.logical_and(dsVar15_2["uniqName"] > 0, dsVar15_2["uniqName"] < 1))
dsVar20_2["uniqName"] = dsVar20_2["uniqName"].where(np.logical_and(dsVar20_2["uniqName"] > 0, dsVar20_2["uniqName"] < 1))
else:
dsVar15_1 = open_files('ncar', SecName.upper()+'-PK', TabName, GeoPK15_case1, Vname_NCAR)
dsVar20_1 = open_files('ncar', SecName.upper()+'-PK', TabName, GeoPK20_case1, Vname_NCAR)
dsVar15_2 = open_files('ncar', SecName.upper()+'-PK', TabName, GeoPK15_case2, Vname_NCAR)
dsVar20_2 = open_files('ncar', SecName.upper()+'-PK', TabName, GeoPK20_case2, Vname_NCAR)
## ESGF
if any (e in Vname_ESGF for e in opstr):
Sym_Expr_ESGF = sympify(Vname_ESGF)
FreeSymb_ESGF = list(Sym_Expr_ESGF.free_symbols)
fESGF = lambdify(FreeSymb_ESGF, Sym_Expr_ESGF, "numpy")
dsTmpbs = []
for s in FreeSymb_ESGF:
dsTmpbs.append(open_files('esgf', SecName.lower()+'-over', TabName, '', s.name))
dsVarbs = dsTmpbs[0].copy()
dsVarbs["uniqName"] = fESGF(*[ds[s.name] for s, ds in zip(FreeSymb_ESGF,dsTmpbs)])
if not noLimit:
dsVarbs["uniqName"] = dsVarbs["uniqName"].where(np.logical_and(dsVarbs["uniqName"] > 0, dsVarbs["uniqName"] < 1))
else:
if Vname_ESGF == 'tas':
dsVarbs = open_files('esgf', SecName.lower()+'-over', 'Amon', '', Vname_ESGF)
else:
dsVarbs = open_files('esgf', SecName.lower()+'-over', TabName, '', Vname_ESGF)
if Vname_E3SM:
if any (e in Vname_E3SM for e in opstr):
Sym_Expr_E3SM = sympify(Vname_E3SM)
FreeSymb_E3SM = list(Sym_Expr_E3SM.free_symbols)
fE3SM = lambdify(FreeSymb_E3SM, Sym_Expr_E3SM, "numpy")
dsTmpbs = []
dsTmp20 = []
for s in FreeSymb_E3SM:
dsTmp20.append(open_files('e3sm', SecName.lower()+'-over', TabName, E3SM_Geo_PK20, s.name))
dsTmpbs.append(open_files('e3sm', SecName.lower()+'-over', TabName, E3SM_Geo_534S, s.name))
dsVar20e = dsTmp20[0].copy()
dsVarbse = dsTmpbs[0].copy()
dsVar20e["uniqName"] = fE3SM(*[ds[s.name] for s, ds in zip(FreeSymb_E3SM,dsTmp20)])
dsVarbse["uniqName"] = fE3SM(*[ds[s.name] for s, ds in zip(FreeSymb_E3SM,dsTmpbs)])
if not noLimit:
dsVar20e["uniqName"] = dsVar20e["uniqName"].where(np.logical_and(dsVar20e["uniqName"] > 0, dsVar20e["uniqName"] < 1))
dsVarbse["uniqName"] = dsVarbse["uniqName"].where(np.logical_and(dsVarbse["uniqName"] > 0, dsVarbse["uniqName"] < 1))
else:
dsVar20e = open_files('e3sm', SecName.lower()+'-over', TabName, E3SM_Geo_PK20, Vname_E3SM)
dsVarbse = open_files('e3sm', SecName.lower()+'-over', TabName, E3SM_Geo_534S, Vname_E3SM)
#mxu read in the cmip6 areacella ana sftlt
areacella = open_files('esgf', SecName.lower()+'-over', 'fx', '', 'areacella')*1.e-6
sftlf = open_files('esgf', SecName.lower()+'-over', 'fx', '', 'sftlf') *1.e-2
if any (e in Vname_NCAR for e in opstr):
if TabName == 'Lmon':
Var15_1_TsGBL = gm_by_area_wght(dsVar15_1, 'uniqName', OpType, True, None, None, TabName) * Sca_NCAR + Fct_NCAR # gC/m2 to kgC/m2
Var20_1_TsGBL = gm_by_area_wght(dsVar20_1, 'uniqName', OpType, True, None, None, TabName) * Sca_NCAR + Fct_NCAR
Var15_2_TsGBL = gm_by_area_wght(dsVar15_2, 'uniqName', OpType, True, None, None, TabName) * Sca_NCAR + Fct_NCAR # gC/m2 to kgC/m2
Var20_2_TsGBL = gm_by_area_wght(dsVar20_2, 'uniqName', OpType, True, None, None, TabName) * Sca_NCAR + Fct_NCAR
else:
Var15_1_TsGBL = gm_by_area_wght(dsVar15_1, 'uniqName', OpType, False, areacella, sftlf, TabName) * Sca_NCAR + Fct_NCAR # gC/m2 to kgC/m2
Var20_1_TsGBL = gm_by_area_wght(dsVar20_1, 'uniqName', OpType, False, areacella, sftlf, TabName) * Sca_NCAR + Fct_NCAR
Var15_2_TsGBL = gm_by_area_wght(dsVar15_2, 'uniqName', OpType, False, areacella, sftlf, TabName) * Sca_NCAR + Fct_NCAR # gC/m2 to kgC/m2
Var20_2_TsGBL = gm_by_area_wght(dsVar20_2, 'uniqName', OpType, False, areacella, sftlf, TabName) * Sca_NCAR + Fct_NCAR
else:
if TabName == 'Lmon':
Var15_1_TsGBL = gm_by_area_wght(dsVar15_1, Vname_NCAR, OpType, True, None, None, TabName) * Sca_NCAR + Fct_NCAR # gC/m2 to kgC/m2
Var20_1_TsGBL = gm_by_area_wght(dsVar20_1, Vname_NCAR, OpType, True, None, None, TabName) * Sca_NCAR + Fct_NCAR
Var15_2_TsGBL = gm_by_area_wght(dsVar15_2, Vname_NCAR, OpType, True, None, None, TabName) * Sca_NCAR + Fct_NCAR # gC/m2 to kgC/m2
Var20_2_TsGBL = gm_by_area_wght(dsVar20_2, Vname_NCAR, OpType, True, None, None, TabName) * Sca_NCAR + Fct_NCAR
else:
Var15_1_TsGBL = gm_by_area_wght(dsVar15_1, Vname_NCAR, OpType, False, areacella, sftlf, TabName) * Sca_NCAR + Fct_NCAR # gC/m2 to kgC/m2
Var20_1_TsGBL = gm_by_area_wght(dsVar20_1, Vname_NCAR, OpType, False, areacella, sftlf, TabName) * Sca_NCAR + Fct_NCAR
Var15_2_TsGBL = gm_by_area_wght(dsVar15_2, Vname_NCAR, OpType, False, areacella, sftlf, TabName) * Sca_NCAR + Fct_NCAR # gC/m2 to kgC/m2
Var20_2_TsGBL = gm_by_area_wght(dsVar20_2, Vname_NCAR, OpType, False, areacella, sftlf, TabName) * Sca_NCAR + Fct_NCAR
if any (e in Vname_ESGF for e in opstr):
Varbs_TsGBL = gm_by_area_wght(dsVarbs, 'uniqName', OpType, False, areacella, sftlf, TabName, TimeShift=False) * Sca_ESGF + Fct_ESGF
else:
Varbs_TsGBL = gm_by_area_wght(dsVarbs, Vname_ESGF, OpType, False, areacella, sftlf, TabName, TimeShift=False) * Sca_ESGF + Fct_ESGF
if Vname_E3SM:
E3SM_areacella = open_files('e3sm', SecName.lower()+'-over', 'fx', E3SM_Geo_HIST, 'areacella')*1.e-6
E3SM_sftlf = open_files('e3sm', SecName.lower()+'-over', 'fx', E3SM_Geo_HIST, 'sftlf')*1.e-2
if any (e in Vname_E3SM for e in opstr):
Varbse_TsGBL = gm_by_area_wght(dsVarbse, 'uniqName', OpType, False, E3SM_areacella, E3SM_sftlf, TabName, TimeShift=True) * Sca_E3SM + Fct_E3SM
Var20e_TsGBL = gm_by_area_wght(dsVar20e, 'uniqName', OpType, False, E3SM_areacella, E3SM_sftlf, TabName, TimeShift=True) * Sca_E3SM + Fct_E3SM
else:
Varbse_TsGBL = gm_by_area_wght(dsVarbse, Vname_E3SM, OpType, False, E3SM_areacella, E3SM_sftlf, TabName, TimeShift=True) * Sca_E3SM + Fct_E3SM
Var20e_TsGBL = gm_by_area_wght(dsVar20e, Vname_E3SM, OpType, False, E3SM_areacella, E3SM_sftlf, TabName, TimeShift=True) * Sca_E3SM + Fct_E3SM
if Vname_E3SM:
if not SkipPlot:
plot_ts_diff([Var15_1_TsGBL, Var15_2_TsGBL, Var20_1_TsGBL, Var20_2_TsGBL, Varbs_TsGBL], Cumdiff, ['CESM2W6-Geo-1.5-1', 'CESM2W6-Geo-1.5-2', 'CESM2W6-Geo-2.0-1', 'CESM2W6-Geo-2.0-2', 'CESM2W6-SSP534OS'], Ylabel,
VarE3SM=[Var20e_TsGBL, Varbse_TsGBL])
return [Var15_1_TsGBL, Var15_2_TsGBL, Var20_1_TsGBL, Var20_2_TsGBL, Varbs_TsGBL, Var20e_TsGBL, Varbse_TsGBL]
else:
if not SkipPlot:
plot_ts_diff([Var15_1_TsGBL, Var15_2_TsGBL, Var20_1_TsGBL, Var20_2_TsGBL, Varbs_TsGBL], Cumdiff, ['CESM2W6-Geo-1.5-1', 'CESM2W6-Geo-1.5-2', 'CESM2W6-Geo-2.0-1', 'CESM2W6-Geo-2.0-2', 'CESM2W6-SSP534OS'], Ylabel)
return [Var15_1_TsGBL, Var15_2_TsGBL, Var20_1_TsGBL, Var20_2_TsGBL, Varbs_TsGBL]
